# %% 导入库
import tensorflow as tf
import matplotlib.pyplot as plt
import os
from tensorflow.keras import layers
import time
from IPython import display
import os
from PIL import Image
import pathlib
import numpy as np

EPOCHS = 200 # 训练轮次
noise_dim = 160 # 噪点大小
num_examples_to_generate = 9 # 生成样例个数
BATCH_SIZE = 256 # 批次大小

# %% 加载数据 将图片为文件转为数组
dir_path = "face"
data_dir = pathlib.Path(dir_path)
imgs = list(data_dir.glob('*.jpg'))
img_arr = []
for img in imgs:
    img = Image.open(str(img))
    img_arr.append(np.array(img)) 
train_images = np.array(img_arr)
nums = train_images.shape[0]
train_images = train_images.reshape(nums, 56, 56, 1).astype('float32')
# 归一化
train_images = (train_images - 127.5) / 127.5  
# 转为tensor格式
train_dataset = tf.data.Dataset.from_tensor_slices(train_images).shuffle(nums).batch(BATCH_SIZE)

# %% 生成器
def make_generator_model():
    model = tf.keras.Sequential()
    model.add(layers.Dense(7*7*256, use_bias=False, input_shape=(noise_dim,)))
    model.add(layers.BatchNormalization())
    model.add(layers.LeakyReLU())

    model.add(layers.Reshape((7, 7, 256)))
    assert model.output_shape == (None, 7, 7, 256)  # Note: None is the batch size

    model.add(layers.Conv2DTranspose(128, (5, 5), strides=(1, 1), padding='same', use_bias=False))
    assert model.output_shape == (None, 7, 7, 128)
    model.add(layers.BatchNormalization())
    model.add(layers.LeakyReLU())

    model.add(layers.Conv2DTranspose(64, (5, 5), strides=(2, 2), padding='same', use_bias=False))
    assert model.output_shape == (None, 14, 14, 64)
    model.add(layers.BatchNormalization())
    model.add(layers.LeakyReLU())

    model.add(layers.Conv2DTranspose(32, (5, 5), strides=(2, 2), padding='same', use_bias=False))
    assert model.output_shape == (None, 28, 28, 32)
    model.add(layers.BatchNormalization())
    model.add(layers.LeakyReLU())

    model.add(layers.Conv2DTranspose(1, (5, 5), strides=(2, 2), padding='same', use_bias=False, activation='tanh'))
    assert model.output_shape == (None, 56, 56, 1)

    return model
   
generator = make_generator_model()
noise = tf.random.normal([1, noise_dim])
# % 试验生成器
# generated_image = generator(noise, training=False)
# plt.imshow(generated_image[0, :, :, 0], cmap='gray')

# %% 鉴别器
def make_discriminator_model():
    model = tf.keras.Sequential()
    model.add(layers.Conv2D(64, (5, 5), strides=(2, 2), padding='same', input_shape=[56, 56, 1]))
    model.add(layers.LeakyReLU())
    model.add(layers.Dropout(0.3))

    model.add(layers.Conv2D(128, (5, 5), strides=(2, 2), padding='same'))
    model.add(layers.LeakyReLU())
    model.add(layers.Dropout(0.3))

    model.add(layers.Flatten())
    model.add(layers.Dense(1))

    return model

discriminator = make_discriminator_model()
# %试验鉴别器
# decision = discriminator(generated_image)
# print (decision)

# %% 配置模型
cross_entropy = tf.keras.losses.BinaryCrossentropy(from_logits=True)

def discriminator_loss(real_output, fake_output):
    real_loss = cross_entropy(tf.ones_like(real_output), real_output)
    fake_loss = cross_entropy(tf.zeros_like(fake_output), fake_output)
    total_loss = real_loss + fake_loss
    return total_loss

def generator_loss(fake_output):
    return cross_entropy(tf.ones_like(fake_output), fake_output)

generator_optimizer = tf.keras.optimizers.Adam(1e-4)
discriminator_optimizer = tf.keras.optimizers.Adam(1e-4)

checkpoint_dir = './training_checkpoints'
checkpoint_prefix = os.path.join(checkpoint_dir, "ckpt")
checkpoint = tf.train.Checkpoint(generator_optimizer=generator_optimizer,
                                 discriminator_optimizer=discriminator_optimizer,
                                 generator=generator,
                                 discriminator=discriminator)

if os.path.exists(checkpoint_dir+"/checkpoint"):
  checkpoint.restore(tf.train.latest_checkpoint(checkpoint_dir))

# 生成样例
seed = tf.random.normal([num_examples_to_generate, noise_dim])
def generate_and_save_images(model, epoch, test_input):
  predictions = model(test_input, training=False)
  plt.figure(figsize=(4, 4))
  for i in range(predictions.shape[0]):
      plt.subplot(3, 3, i+1)
      plt.imshow(predictions[i, :, :, 0] * 127.5 + 127.5, cmap='gray')
      plt.axis('off')
  plt.savefig('image_at_epoch_{:04d}.png'.format(epoch))
  plt.show()

# %% 训练
@tf.function
def train_step(images):
    noise = tf.random.normal([BATCH_SIZE, noise_dim])

    with tf.GradientTape() as gen_tape, tf.GradientTape() as disc_tape:
      generated_images = generator(noise, training=True)

      real_output = discriminator(images, training=True)
      fake_output = discriminator(generated_images, training=True)

      gen_loss = generator_loss(fake_output)
      disc_loss = discriminator_loss(real_output, fake_output)

    gradients_of_generator = gen_tape.gradient(gen_loss, generator.trainable_variables)
    gradients_of_discriminator = disc_tape.gradient(disc_loss, discriminator.trainable_variables)

    generator_optimizer.apply_gradients(zip(gradients_of_generator, generator.trainable_variables))
    discriminator_optimizer.apply_gradients(zip(gradients_of_discriminator, discriminator.trainable_variables))

# 训练
def train(dataset, epochs):
  for epoch in range(epochs):
    start = time.time()
    for image_batch in dataset:
      train_step(image_batch)

    display.clear_output(wait=True)
    generate_and_save_images(generator, epoch + 1, seed)
    if (epoch + 1) % 20 == 0: # 20保存一次
      checkpoint.save(file_prefix = checkpoint_prefix)
    print ('Time for epoch {} is {} sec'.format(epoch + 1, time.time()-start))

#%% 生成一个
def prediction():
  test_input = tf.random.normal([1, noise_dim])
  predictions = generator(test_input, training=False)
  img_arr = predictions[0][:, :, 0]
  img_arr = img_arr* 127.5 + 127.5
  plt.axis('off')
  plt.figure(figsize=(1, 1))
  plt.imshow(img_arr, cmap='gray')
  plt.show()
  
  # 多个
  seed = tf.random.normal([num_examples_to_generate, noise_dim])
  generate_and_save_images(generator, 1, seed)

# %%
if __name__ == "__main__":
  print("welcome")
  # 训练
  # train(train_dataset, EPOCHS)
  # 生成
  # prediction()
# %%
